package cc.autoapi.pucong.auto.flow.provider.api.ai;


import cc.autoapi.pucong.auto.flow.core.flow.execNode.entity.node.*;
import cn.hutool.core.lang.UUID;
import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONArray;
import com.alibaba.fastjson2.JSONObject;
import com.alibaba.fastjson2.JSONWriter;
import cc.autoapi.pucong.auto.flow.core.flow.execNode.ai.send.AbstractAiSendSendTool;
import cc.autoapi.pucong.auto.flow.core.flow.execNode.ai.send.handler.TransAiHandlerSend;
import cc.autoapi.pucong.auto.flow.core.flow.execNode.selfenum.DataTypeEnum;
import cc.autoapi.pucong.auto.flow.core.flow.execNode.selfenum.NodeTypeEnum;
import cc.autoapi.pucong.auto.flow.core.flow.execNode.task.AbstractTask;
import cc.autoapi.pucong.auto.flow.core.flow.execNode.task.TaskFactory;
import cc.autoapi.pucong.auto.flow.core.flow.execNode.utils.AiJsonUtils;

import java.util.ArrayList;
import java.util.List;
import java.util.Set;

/**
 * @description:
 * @author: cong.pu
 * @create: 2024-10-15 09:08
 **/
public class AiFlowTest {

    private static String kimiKey = "";
    private static String tran_appid = "";
    private static String TRAN_APPKEY = "";
    private static String prompt = "  你是一个流程编排大师。根据提供组件列表、画布中已存在的组件实例、已存在的组件实例关系和功能描述，调用工具完成流程编排。\n" +
            "                        \n" +
            "            你的任务是：\n" +
            "            1. 在画布中创建组件实例，并设置组件实例之间的关系和属性。\n" +
            "            2. 根据组件的规则，某些组件可以添加输入和输出参数，你需要为这些特殊组件添加输入和输出参数。\n" +
            "            3. 设置组件的输入值来源，可以是静态值、前面组件的输出参数、或全局参数。\n" +
            "            4. 使用工具构建清晰的规则和逻辑结构，确保每个组件实例至少有一个关系。\n" +
            "            5. 当调用工具返回执行失败时，分析原因并修正参数，再次调用工具。\n" +
            "                        \n" +
            "            流程编排的规则：\n" +
            "            1. 组件编码为“apiConfig”、“cronConfig”的组件是起始节点。\n" +
            "            2. 先执行起始组件，根据关系决定执行顺序。\n" +
            "            3. 当执行节点时，判断上一节点是否执行完成，上一节点都执行完成才能执行本节点。当前组件执行完成后，下一执行组件有多个时，异步同时执行下一节点。\n" +
            "            4. 节点之间的关系类型有：0：一般执行关系，1：catch执行关系，2：工具执行关系，3：true执行关系，4：false执行关系。\n" +
            "            5. 组件jsScript和python只做数据处理，复杂功能需交给java脚本组件。\n" +
            "            6. 输入输出参数的数据类型为[\"number\", \"string\", \"txt\", \"Boolean\", \"Date\", \"List\", \"Object\"]。\n" +
            "            7. 组件实例下的输入参数编码不能相同，输出参数编码也不能相同。\n" +
            "            8. 添加了实例组件后，需要添加关系，才能生成正确的流程图。\n" +
            "                        \n" +
            "            组件列表：\n" +
            "            {nodes}\n" +
            "                        \n" +
            "            画布中的组件：\n" +
            "            {nodeInstances}\n" +
            "                        \n" +
            "            画布中的组件关系：\n" +
            "            {nodeInsRelations}\n" +
            "                        \n" +
            "            功能描述：\n" +
            "            {desc}";
    private static JSONArray messages = new JSONArray();
    private static List<TaskNode> taskNodeList = new ArrayList<>();

    private static List<TaskNodeRelation> taskNodeRelationList = new ArrayList<>();

    //    private static AbstractAiSendSendTool aiSendSendTool = new Kimi32kAiHandlerSend(kimiKey);
    private static AbstractAiSendSendTool aiSendSendTool = new TransAiHandlerSend(tran_appid, TRAN_APPKEY);

    public static void main(String[] args) {
        // 获取组件列表
        String nodes = getNodes();
        // 创建工具
        JSONArray tools = AiJsonUtils.getToolsJon();
        aiSendSendTool.initTools(tools);
        String nodeInstances = getNodeInstances();
        String nodeInsRelations = getNodeInsRelations();
        // 发送Ai
        String system = prompt.replace("{nodes}", nodes);
        system = system.replace("{nodeInstances}", nodeInstances);
        system = system.replace("{nodeInsRelations}", nodeInsRelations);
        system = system.replace("{desc}", getFirstNodeInfo().getString("描述"));

        JSONObject startNodeInfo = new JSONObject();
        startNodeInfo.put("role", "user");
        startNodeInfo.put("content", system);
        messages.add(startNodeInfo);

        System.out.println(JSON.toJSONString(messages, JSONWriter.Feature.PrettyFormat));
        JSONObject jsonObject = aiSendSendTool.sendToolsToAi(messages);
        // 使用工具
        userTool(jsonObject);
        // 输出设计内容
        System.out.println(JSON.toJSONString(taskNodeList, JSONWriter.Feature.PrettyFormat));
        System.out.println(JSON.toJSONString(taskNodeRelationList, JSONWriter.Feature.PrettyFormat));
    }


    private static String getNodeInsRelations() {
        JSONArray relations = new JSONArray();
        for (TaskNodeRelation nodeRelation : taskNodeRelationList) {
            JSONObject jsonObject = new JSONObject();
            jsonObject.put("来源实例组件id", nodeRelation.getNodeFromId());
            jsonObject.put("目标实例组件id", nodeRelation.getNodeToId());
            jsonObject.put("关系类型", nodeRelation.getRelationType());
            relations.add(jsonObject);
        }
        return relations.toJSONString();
    }

    private static String getNodeInstances() {
        // 添加首节点
        TaskNode firstNode = new TaskNode();
        firstNode.setUuid(UUID.randomUUID().toString());
        firstNode.setNodeName("智能搜索");
        firstNode.setNodeType(NodeTypeEnum.API_CONFIG.getKey());
        firstNode.setNodeDesc("调用百度网站的搜索接口，并获取所有的标题");

        List<TaskNodeParamIn> ins = new ArrayList<>();
        TaskNodeParamIn in2 = new TaskNodeParamIn();
        in2.setUuid(UUID.randomUUID().toString());
        in2.setParamCode("keywords");
        in2.setParamDesc("需要搜索的内容，关键词");
        in2.setDateType(DataTypeEnum.STRING);
        ins.add(in2);
        List<TaskNodeParamOut> outs = new ArrayList<>();
        TaskNodeParamOut out2 = new TaskNodeParamOut();
        out2.setUuid(UUID.randomUUID().toString());
        out2.setParamCode("titleList");
        out2.setParamDesc("搜索到的标题列表");
        out2.setDateType(DataTypeEnum.LIST);
        outs.add(out2);
        firstNode.setTaskNodeParamInList(ins);
        firstNode.setTaskNodeParamOutList(outs);
        taskNodeList.add(firstNode);
        JSONArray nodes = new JSONArray();
        for (TaskNode taskNode : taskNodeList) {
            JSONObject jsonObject = new JSONObject();
            jsonObject.put("实例组件id", taskNode.getUuid());
            jsonObject.put("组件名称", taskNode.getNodeName());
            jsonObject.put("组件编码", taskNode.getNodeType());
            jsonObject.put("描述", taskNode.getNodeDesc());

            JSONArray paramsInList = new JSONArray();
            List<TaskNodeParamIn> taskNodeParamInList = taskNode.getTaskNodeParamInList();
            if (taskNodeParamInList != null) {
                for (TaskNodeParamIn in : taskNodeParamInList) {
                    JSONObject apiNodeParamsInVo = new JSONObject();
                    apiNodeParamsInVo.put("参数编码", in.getParamCode());
                    apiNodeParamsInVo.put("数据类型", in.getDateType().getName());
                    apiNodeParamsInVo.put("是否必填", in.getNeedValue());
                    apiNodeParamsInVo.put("描述", in.getParamDesc());
                    paramsInList.add(apiNodeParamsInVo);
                }
            }
            JSONArray paramsOutList = new JSONArray();
            List<TaskNodeParamOut> taskNodeParamOutList = taskNode.getTaskNodeParamOutList();
            if (taskNodeParamOutList != null) {
                for (TaskNodeParamOut out : taskNodeParamOutList) {
                    JSONObject apiNodeParamsOutVo = new JSONObject();
                    apiNodeParamsOutVo.put("参数编码", out.getParamCode());
                    apiNodeParamsOutVo.put("数据类型", out.getDateType().getName());
                    apiNodeParamsOutVo.put("描述", out.getParamDesc());
                    paramsOutList.add(apiNodeParamsOutVo);
                }
            }
            jsonObject.put("输入参数列表", paramsInList.toJSONString());
            jsonObject.put("输出参数列表", paramsOutList.toJSONString());
            nodes.add(jsonObject);
        }
        return nodes.toJSONString();
    }

    private static JSONObject getFirstNodeInfo() {
        TaskNode taskNode = taskNodeList.get(0);
        JSONObject base1 = new JSONObject();
        base1.put("实例组件id", taskNode.getUuid());
        base1.put("组件名称", taskNode.getNodeName());
        base1.put("组件编码", taskNode.getNodeType());
        base1.put("描述", taskNode.getNodeDesc());

        List<TaskNodeParamIn> propertyParamsIns = taskNode.getTaskNodeParamInList();
        JSONArray paramsInList = new JSONArray();
        for (TaskNodeParamIn nodeParamIn : propertyParamsIns) {
            JSONObject apiNodeParamsInVo = new JSONObject();
            apiNodeParamsInVo.put("主键id", nodeParamIn.getUuid());
            apiNodeParamsInVo.put("参数编码", nodeParamIn.getParamCode());
            apiNodeParamsInVo.put("数据类型", nodeParamIn.getDateType().getType());
            apiNodeParamsInVo.put("是否必填", nodeParamIn.getNeedValue());
            apiNodeParamsInVo.put("描述", nodeParamIn.getParamDesc());
            paramsInList.add(apiNodeParamsInVo);
        }
        base1.put("输入参数列表", paramsInList);
        JSONArray paramsOutList = new JSONArray();
        List<TaskNodeParamOut> propertyParamsOuts = taskNode.getTaskNodeParamOutList();
        for (TaskNodeParamOut nodeParamOut : propertyParamsOuts) {
            JSONObject apiNodeParamsOutVo = new JSONObject();
            apiNodeParamsOutVo.put("主键id", nodeParamOut.getUuid());
            apiNodeParamsOutVo.put("参数编码", nodeParamOut.getParamCode());
            apiNodeParamsOutVo.put("数据类型", nodeParamOut.getDateType().getType());
            apiNodeParamsOutVo.put("描述", nodeParamOut.getParamDesc());
            paramsOutList.add(apiNodeParamsOutVo);
        }
        base1.put("输出参数列表", paramsOutList);

        return base1;
    }


    private static String getNodes() {
        JSONArray nodes = new JSONArray();
        NodeTypeEnum[] values = NodeTypeEnum.values();
        for (NodeTypeEnum value : values) {
            if (value == NodeTypeEnum.API_CONFIG || value == NodeTypeEnum.CRON_CONFIG) {
                continue;
            }
            AbstractTask task = TaskFactory.getTask(value);
            JSONObject node = getNode(value.getKey(), task);
            nodes.add(node);
        }

        Set<String> taskMapKey = TaskFactory.getTaskMapKey();
        for (String s : taskMapKey) {
            AbstractTask task = TaskFactory.getTask(s);
            JSONObject node = getNode(s, task);
            nodes.add(node);
        }
        return nodes.toJSONString();
    }

    private static JSONObject getNode(String nodeCode, AbstractTask task) {
        JSONObject base1 = new JSONObject();
        base1.put("组件名称", task.getTaskName());
        base1.put("组件编码", nodeCode);
        base1.put("描述", task.getTaskDesc());
        base1.put("自定义输入和输出参数", task.getTaskParamInAuto());
        List<NodeParamIn> propertyParamsIns = task.getPropertyParamsIns();
        JSONArray paramsInList = new JSONArray();
        for (NodeParamIn nodeParamIn : propertyParamsIns) {
            JSONObject apiNodeParamsInVo = new JSONObject();
            apiNodeParamsInVo.put("参数编码", nodeParamIn.getParamCode());
            apiNodeParamsInVo.put("数据类型", nodeParamIn.getDateType().getType());
            apiNodeParamsInVo.put("是否必填", nodeParamIn.getNeedValue());
            apiNodeParamsInVo.put("描述", nodeParamIn.getDesc());
            paramsInList.add(apiNodeParamsInVo);
        }
        base1.put("输入参数列表", paramsInList);
        JSONArray paramsOutList = new JSONArray();
        List<NodeParamOut> propertyParamsOuts = task.getPropertyParamsOuts();
        for (NodeParamOut nodeParamOut : propertyParamsOuts) {
            JSONObject apiNodeParamsOutVo = new JSONObject();
            apiNodeParamsOutVo.put("参数编码", nodeParamOut.getParamCode());
            apiNodeParamsOutVo.put("数据类型", nodeParamOut.getDateType().getType());
            apiNodeParamsOutVo.put("描述", nodeParamOut.getDesc());
            paramsOutList.add(apiNodeParamsOutVo);
        }
        base1.put("输出参数列表", paramsOutList);
        return base1;
    }

    private static String getProperties() {
        JSONArray properties = new JSONArray();

        JSONObject property = new JSONObject();
        property.put("参数编码", "authCode");
        property.put("数据类型", DataTypeEnum.STRING);
        property.put("描述", "授权码，访问网站时使用");
        properties.add(property);

        return properties.toJSONString();
    }

    private static void userTool(JSONObject jsonObject) {
        if (jsonObject != null) {
            messages.add(jsonObject);
            JSONArray toolCalls = jsonObject.getJSONArray("tool_calls");
            if (toolCalls != null) {
                for (Object toolCall : toolCalls) {
                    JSONObject toolCallJson = (JSONObject) toolCall;
                    String id = toolCallJson.getString("id");
                    String toolCallName = toolCallJson.getJSONObject("function").getString("name");
                    JSONArray toolCallArguments = toolCallJson.getJSONObject("function").getJSONArray("arguments");
                    // 执行工具
                    JSONObject toolsResult = runResult(toolCallName, toolCallArguments);
                    // 设置结果
                    JSONObject msg = new JSONObject();
                    msg.put("role", "tool");
                    msg.put("tool_call_id", id);
                    msg.put("name", toolCallName);
                    msg.put("content", JSON.toJSONString(toolsResult));
                    messages.add(msg);
                }
                // 再次调用
                System.out.println(JSON.toJSONString(messages, JSONWriter.Feature.PrettyFormat));
                JSONObject jsonObject1 = aiSendSendTool.sendToolsToAi(messages);
                userTool(jsonObject1);
            }
        }
    }

    private static JSONObject runResult(String toolCallName, JSONArray toolCallArguments) {
        System.out.println("执行工具：" + toolCallName);
        System.out.println("参数：" + toolCallArguments.toJSONString());
        JSONObject result = new JSONObject();
        result.put("执行状态", true);

        JSONObject jsonObject = toolCallArguments.getJSONObject(0);
        if ("addNode".equals(toolCallName)) {
            String nodeName = jsonObject.getString("nodeName");
            String nodeCode = jsonObject.getString("nodeCode");
            String nodeDesc = jsonObject.getString("nodeDesc");
            TaskNode node = new TaskNode();
            node.setUuid(UUID.randomUUID().toString());
            node.setNodeName(nodeName);
            node.setNodeType(nodeCode);
            node.setNodeDesc(nodeDesc);
            taskNodeList.add(node);
            result.put("实例组件id", node.getUuid());
        } else if ("addNodeRelation".equals(toolCallName)) {
            Integer relationType = jsonObject.getInteger("relationType");
            String fromId = jsonObject.getString("fromId");
            String nodeToId = jsonObject.getString("nodeToId");
            TaskNode from = taskNodeList.stream().filter(t -> t.getUuid().equals(fromId)).findFirst().orElse(null);
            if (from == null) {
                result.put("执行状态", false);
                result.put("失败原因", "根据来源组件实例id，未找到关系目标组件实例");
                return result;
            }
            TaskNode to = taskNodeList.stream().filter(t -> t.getUuid().equals(nodeToId)).findFirst().orElse(null);
            if (to == null) {
                result.put("执行状态", false);
                result.put("失败原因", "根据目标组件实例id，未找到关系目标组件实例");
                return result;
            }
            TaskNodeRelation relation = new TaskNodeRelation(relationType, fromId, nodeToId);
            taskNodeRelationList.add(relation);
        } else if ("addParamIn".equals(toolCallName)) {
            String instanceNodeId = jsonObject.getString("instanceNodeId");
            String paramCode = jsonObject.getString("paramCode");
            String dataType = jsonObject.getString("dataType");
            String paramDesc = jsonObject.getString("paramDesc");
            TaskNode taskNode = taskNodeList.stream().filter(t -> t.getUuid().equals(instanceNodeId)).findFirst().orElse(null);
            if (taskNode == null) {
                result.put("执行状态", false);
                result.put("失败原因", "根据组件实例id，未找到组件实例");
                return result;
            }
            DataTypeEnum itemByName = DataTypeEnum.getItemByName(dataType);
            if (itemByName == null) {
                result.put("执行状态", false);
                result.put("失败原因", "设置的数据类型不对，没有对应的枚举类型：" + dataType);
                return result;
            }
            List<TaskNodeParamIn> taskNodeParamInList = taskNode.getTaskNodeParamInList();
            if (taskNodeParamInList == null) {
                taskNodeParamInList = new ArrayList<>();
                taskNode.setTaskNodeParamInList(taskNodeParamInList);
            }
            long count = taskNodeParamInList.stream().filter(t -> t.getParamCode().equals(paramCode)).count();
            if (count > 0) {
                result.put("执行状态", false);
                result.put("失败原因", "添加的输入参数编码已存在，同一组件下输入参数编码不能相同！");
                return result;
            }
            TaskNodeParamIn in = new TaskNodeParamIn();
            in.setUuid(UUID.randomUUID().toString());
            in.setParamCode(paramCode);
            in.setParamDesc(paramDesc);
            in.setDateType(itemByName);
            taskNodeParamInList.add(in);
            result.put("组件实例输入参数id", in.getUuid());
        } else if ("addParamOut".equals(toolCallName)) {
            String instanceNodeId = jsonObject.getString("instanceNodeId");
            String paramCode = jsonObject.getString("paramCode");
            String dataType = jsonObject.getString("dataType");
            String paramDesc = jsonObject.getString("paramDesc");
            Boolean returnValue = jsonObject.getBoolean("returnValue");
            TaskNode taskNode = taskNodeList.stream().filter(t -> t.getUuid().equals(instanceNodeId)).findFirst().orElse(null);
            if (taskNode == null) {
                result.put("执行状态", false);
                result.put("失败原因", "根据组件实例id，未找到组件实例");
                return result;
            }
            DataTypeEnum itemByName = DataTypeEnum.getItemByName(dataType);
            if (itemByName == null) {
                result.put("执行状态", false);
                result.put("失败原因", "设置的数据类型不对，没有对应的枚举类型：" + dataType);
                return result;
            }
            List<TaskNodeParamOut> taskNodeParamOutList = taskNode.getTaskNodeParamOutList();
            if (taskNodeParamOutList == null) {
                taskNodeParamOutList = new ArrayList<>();
                taskNode.setTaskNodeParamOutList(taskNodeParamOutList);
            }
            long count = taskNodeParamOutList.stream().filter(t -> t.getParamCode().equals(paramCode)).count();
            if (count > 0) {
                result.put("执行状态", false);
                result.put("失败原因", "添加的输出参数编码已存在，同一组件下输出参数编码不能相同！");
                return result;
            }
            TaskNodeParamOut out = new TaskNodeParamOut();
            out.setUuid(UUID.randomUUID().toString());
            out.setParamCode(paramCode);
            out.setParamDesc(paramDesc);
            out.setDateType(itemByName);
            out.setReturnValue(returnValue);
            taskNodeParamOutList.add(out);
            result.put("组件实例输出参数id", out.getUuid());
        } else if ("setParamInValue".equals(toolCallName)) {
            String instanceNodeId = jsonObject.getString("instanceNodeId");
            String paramCode = jsonObject.getString("paramCode");
            String paramValue = jsonObject.getString("paramValue");
            Integer paramDataType = jsonObject.getInteger("paramDataType");
            TaskNode taskNode = taskNodeList.stream().filter(t -> t.getUuid().equals(instanceNodeId)).findFirst().orElse(null);
            if (taskNode == null) {
                result.put("执行状态", false);
                result.put("失败原因", "根据组件实例id，未找到组件实例");
                return result;
            }
            List<TaskNodeParamIn> taskNodeParamInList = taskNode.getTaskNodeParamInList();
            if (taskNodeParamInList == null) {
                result.put("执行状态", false);
                result.put("失败原因", "根据组件实例id和输入参数code未找到该组件的参数！");
                return result;
            }
            TaskNodeParamIn in = taskNodeParamInList.stream().filter(t -> t.getParamCode().equals(paramCode)).findFirst().orElse(null);
            if (in == null) {
                result.put("执行状态", false);
                result.put("失败原因", "根据组件实例id和输入参数code未找到该组件的参数！");
                return result;
            }
            in.setParamValue(paramValue);
            in.setParamDataType(paramDataType);
        } else if ("setNodeTag".equals(toolCallName)) {
            String instanceNodeId = jsonObject.getString("instanceNodeId");
            String tagName = jsonObject.getString("tagName");
            TaskNode taskNode = taskNodeList.stream().filter(t -> t.getUuid().equals(instanceNodeId)).findFirst().orElse(null);
            if (taskNode == null) {
                result.put("执行状态", false);
                result.put("失败原因", "根据组件实例id，未找到组件实例");
                return result;
            }
            taskNode.setNodeTag(tagName);
        } else if ("setNodeScript".equals(toolCallName)) {
            String instanceNodeId = jsonObject.getString("instanceNodeId");
            String script = jsonObject.getString("script");
            TaskNode taskNode = taskNodeList.stream().filter(t -> t.getUuid().equals(instanceNodeId)).findFirst().orElse(null);
            if (taskNode == null) {
                result.put("执行状态", false);
                result.put("失败原因", "根据组件实例id，未找到组件实例");
                return result;
            }
            taskNode.setNodeScript(script);
        }
        return result;
    }
}
